package algorithm.graph;


import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

public class Prim {

    public static Set<Edge> primMST(Graph graph){
        HashSet<Edge> result = new HashSet<>();
        HashSet<Node> set = new HashSet<>();
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new Comparator<Edge>() {
            @Override
            public int compare(Edge o1, Edge o2) {
                return o1.weight - o2.weight;
            }
        });
        for (Node node : graph.nodes.values()) {
            if(!set.contains(node)) {
                set.add(node);
                for (Edge edge : node.edges) {
                    priorityQueue.add(edge);
                }
                while (!priorityQueue.isEmpty()) {
                    Edge edge = priorityQueue.poll();
                    Node to = edge.to;
                    if(!set.contains(to)){
                        set.add(to);
                        result.add(edge);
                        for (Edge edge1 : to.edges) {
                            priorityQueue.add(edge1);
                        }
                    }
                }
            }

        }
        return result;
    }
}
